import os
from core.writer_utils import KeyPointMsg
from core.writer import Writer

from prompts.创作剧情.prompt import main as prompt_plot
from prompts.提炼.prompt import main as prompt_summary

class PlotWriter(Writer):
    def __init__(self, xy_pairs, global_context, model=None, sub_model=None, x_chunk_length=200, y_chunk_length=1000, max_thread_num=5):
        super().__init__(xy_pairs, global_context, model, sub_model, x_chunk_length=x_chunk_length, y_chunk_length=y_chunk_length, max_thread_num=max_thread_num)

    def write(self, user_prompt, prompt_type, pair_span=None):
        target_chunk = self.get_chunk(pair_span=pair_span)

        if not self.global_context.get("chapter", ''):
            raise Exception("需要提供章节内容。")
        
        if not target_chunk.y_chunk.strip():
            if not self.y.strip():
                chunks = [target_chunk, ]
            else:
                raise Exception("选中进行创作的内容不能为空，考虑随便填写一些占位的字。")
        else:
            chunks = self.get_chunks(pair_span)

        new_chunks = yield from self.batch_yield(
            [self.write_text(e, prompt_plot, user_prompt) for e in chunks], 
            chunks, prompt_name='创作文本', prompt_type=prompt_type)
        
        results = yield from self.batch_map_text(new_chunks)
        new_chunks2 = [e[0] for e in results]

        self.apply_chunks(chunks, new_chunks2)
        
        # 缓存生成的内容到本地文件
        self._cache_generated_content(user_prompt, prompt_type)
    
    def summary(self):
        target_chunk = self.get_chunk(pair_span=(0, len(self.xy_pairs)))
        if not target_chunk.y_chunk:
            raise Exception("没有剧情需要总结。")
        if len(target_chunk.y_chunk) <= 5:
            raise Exception("需要总结的剧情不能少于5个字。")
        
        result = yield from prompt_summary(self.model, "提炼章节", y=target_chunk.y_chunk)

        self.global_context['chapter'] = result['text']

    def get_model(self):
        return self.model

    def get_sub_model(self):
        return self.sub_model
    
    def _cache_generated_content(self, user_prompt, prompt_type):
        """
        将生成的内容缓存到本地文件
        文件名格式: {章节名}_{创作提示词类型}.txt
        保存路径: cache/{文档名}/
        """
        try:
            # 获取章节名
            chapter_name = self.global_context.get("chapter", "unknown_chapter")
            
            # 获取文档名（从全局上下文或通过其他方式）
            novel_name = self.global_context.get("novel_name", "unknown_novel")
            
            # 创建文件名
            filename = f"{chapter_name}_{prompt_type}.txt"
            
            # 确定保存路径
            project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            cache_dir = os.path.join(project_root, "cache", novel_name)
            os.makedirs(cache_dir, exist_ok=True)
            
            # 完整文件路径
            file_path = os.path.join(cache_dir, filename)
            
            # 将生成的内容写入文件
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(self.y)
            
        except Exception as e:
            # 记录错误但不中断主流程
            from core.logger import get_logger
            logger = get_logger(name="Writer")
            logger.warning(f"缓存生成内容时出错: {e}")
